#include <linux/init.h>
#include <linux/module.h>
#include <linux/rbtree.h>

MODULE_LICENSE("Dual BSD/GPL");

typedef struct node_s {
	struct rb_node node;
	__u16 id;
} node_t;

typedef struct tree_s {
	struct rb_root root[1];
	__u16 node_cnt;
} tree_t;

typedef struct test_obj_s {
	tree_t tree[1];
} test_obj_t;

void add_node(tree_t *tree, node_t *new)
{
	struct rb_node **p = NULL;
	struct rb_node *parent = NULL;
	node_t *node_tmp = NULL;	
/*
	if (!tree->root->rb_node) {
		tree->root->rb_node = &new->node;
		rb_set_black(&new->node);
		return;
	}
*/
	p = &(tree->root->rb_node);
	while(*p) {
		parent = *p;
		node_tmp = rb_entry(parent, node_t, node);
		if (new->id < node_tmp->id)
			p = &(*p)->rb_left;
		else if (new->id > node_tmp->id)
			p = &(*p)->rb_right;
		else
			BUG();
	}
	rb_link_node(&new->node, parent, p);
	rb_set_red(&new->node);
	rb_insert_color(&new->node, tree->root);
}

int create_tree(test_obj_t *obj)
{
	//__u16 data[] = {123, 432, 245, 98, 31, 487, 398, 43, 67};
	__u16 data[] = {6, 5, 4, 3, 2, 1};
	int n, i, ret = 0;
	node_t *node = NULL;

	n = sizeof(data) / sizeof(__u16);
	
	for (i = 0; i < n; i++) {
		printk(KERN_ERR "%d \n", data[i]);
		node = kmalloc(sizeof(node_t), GFP_KERNEL);
		if (!node) {
			ret = -ENOMEM;
			break;
		}
		memset(node, 0, sizeof(node_t));
		node->id = data[i];
		add_node(obj->tree, node);
	}
	return ret;
}

void print_test_obj(test_obj_t *obj)
{
	struct rb_node *p;
	node_t *node = NULL;

	p = rb_first(obj->tree->root);
	while(p != NULL) {
		node = rb_entry(p, node_t, node);	
		printk(KERN_ERR "id:%d \n", node->id);
		p = rb_next(p);
	}
}

void tree_destroy(tree_t *tree)
{
	node_t *node = NULL;
	struct rb_node *p = NULL;

	while(tree->root->rb_node) {
		p = rb_first(tree->root);
		node = rb_entry(p, node_t, node);
		rb_erase(&node->node, tree->root);
		kfree(node);
	}
}

void destroy_test_obj(test_obj_t *obj)
{
	tree_destroy(obj->tree);
}

void test_rb_tree(void) {
	test_obj_t *test_obj = NULL;
	int ret = 0;

	test_obj = kmalloc(sizeof(test_obj_t), GFP_KERNEL);
	if (!test_obj) {
		ret = -ENOMEM;
		printk(KERN_ERR "can't alloc memory for test obj \n");
		goto obj_err;
	}
	memset(test_obj, 0, sizeof(test_obj_t));

	ret = create_tree(test_obj);
	if (ret < 0) {
		printk(KERN_ERR "create tree failed\n");
		goto create_err;
	}
	
	print_test_obj(test_obj);
	destroy_test_obj(test_obj);
	print_test_obj(test_obj);

create_err:
	kfree(test_obj);
obj_err:
	return;
}

static __init int test_init(void)
{
	printk(KERN_ERR "test init \n");
	test_rb_tree();
	return 0;
}

static __exit void test_exit(void)
{
	printk(KERN_ERR "test exit \n");
}

module_init(test_init);
module_exit(test_exit);

